#include <stdio.h>
#include <pcap.h>
#include <time.h>
#include <arpa/inet.h>
#include <net/ethernet.h> 
#include <netinet/ip.h>   
#include <string.h>

#define CSV_LOG_FILE "network_traffic_log.csv"
#define STATS_LOG_FILE "network_traffic_stats.csv"
#define INTERVAL 60  // 统计时间间隔（秒）

typedef struct {
    char addr[18];          // MAC 或 IP 地址
    unsigned long tx_len;   // 发送数据长度
    unsigned long rx_len;   // 接收数据长度
} addr_stat;

// 统计 MAC 地址和 IP 地址的流量数据
addr_stat mac_stats[100];
addr_stat ip_stats[100];
int mac_stats_count = 0;
int ip_stats_count = 0;

// 更新统计信息
void update_stats(addr_stat stats[], int* stats_count, const char* src_addr, const char* dst_addr, int length) {
    int found_src = 0, found_dst = 0;

    // 遍历现有记录，更新已有的统计数据
    for (int i = 0; i < *stats_count; i++) {
        if (strcmp(stats[i].addr, src_addr) == 0) {
            stats[i].tx_len += length;
            found_src = 1;
        }
        if (strcmp(stats[i].addr, dst_addr) == 0) {
            stats[i].rx_len += length;
            found_dst = 1;
        }
    }

    // 如果源地址没有记录，新增一条记录
    if (!found_src) {
        strcpy(stats[*stats_count].addr, src_addr);
        stats[*stats_count].tx_len = length;
        stats[*stats_count].rx_len = 0;
        (*stats_count)++;
    }

    // 如果目标地址没有记录，新增一条记录
    if (!found_dst) {
        strcpy(stats[*stats_count].addr, dst_addr);
        stats[*stats_count].tx_len = 0;
        stats[*stats_count].rx_len = length;
        (*stats_count)++;
    }
}

// 记录捕获的数据包到 CSV 文件
void log_packet_to_csv(const char* timestamp, const char* src_mac, const char* src_ip, const char* dst_mac, const char* dst_ip, int length) {
    FILE* log_file = fopen(CSV_LOG_FILE, "a");
    if (log_file == NULL) {
        perror("Error opening CSV log file");
        return;
    }
    fprintf(log_file, "%s,%s,%s,%s,%s,%d\n", timestamp, src_mac, src_ip, dst_mac, dst_ip, length);
    fclose(log_file);
}

// 每隔 60 秒记录 MAC 和 IP 地址的统计信息到 CSV 文件
void log_stats_to_csv() {
    FILE* stats_file = fopen(STATS_LOG_FILE, "a");
    if (stats_file == NULL) {
        perror("Error opening stats log file");
        return;
    }

    time_t now = time(NULL);
    char timestamp[64];
    strftime(timestamp, sizeof timestamp, "%Y-%m-%d %H:%M:%S", localtime(&now));

    // 记录 MAC 地址统计
    fprintf(stats_file, "时间,MAC地址,发送数据长度,接收数据长度\n");
    for (int i = 0; i < mac_stats_count; i++) {
        fprintf(stats_file, "%s,%s,%lu,%lu\n", timestamp, mac_stats[i].addr, mac_stats[i].tx_len, mac_stats[i].rx_len);
    }

    // 记录 IP 地址统计
    fprintf(stats_file, "时间,IP地址,发送数据长度,接收数据长度\n");
    for (int i = 0; i < ip_stats_count; i++) {
        fprintf(stats_file, "%s,%s,%lu,%lu\n", timestamp, ip_stats[i].addr, ip_stats[i].tx_len, ip_stats[i].rx_len);
    }

    fclose(stats_file);
}

// 处理每个捕获的数据包
void packet_handler(u_char* args, const struct pcap_pkthdr* header, const u_char* packet) {
    // 解析以太网头部
    struct ether_header* eth_header = (struct ether_header*)packet;
    char src_mac[18], dst_mac[18];
    snprintf(src_mac, sizeof(src_mac), "%s", ether_ntoa((struct ether_addr*)eth_header->ether_shost));
    snprintf(dst_mac, sizeof(dst_mac), "%s", ether_ntoa((struct ether_addr*)eth_header->ether_dhost));

    // 如果不是 IP 包，跳过
    if (ntohs(eth_header->ether_type) != ETHERTYPE_IP) {
        return;
    }

    // 解析 IP 包
    struct ip* ip_header = (struct ip*)(packet + sizeof(struct ether_header));
    char src_ip[INET_ADDRSTRLEN], dst_ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &(ip_header->ip_src), src_ip, INET_ADDRSTRLEN);
    inet_ntop(AF_INET, &(ip_header->ip_dst), dst_ip, INET_ADDRSTRLEN);

    // 获取时间戳
    time_t now = time(NULL);
    char timestamp[64];
    strftime(timestamp, sizeof timestamp, "%Y-%m-%d %H:%M:%S", localtime(&now));

    // 记录到 CSV
    log_packet_to_csv(timestamp, src_mac, src_ip, dst_mac, dst_ip, header->len);

    // 更新 MAC 和 IP 的统计信息
    update_stats(mac_stats, &mac_stats_count, src_mac, dst_mac, header->len);
    update_stats(ip_stats, &ip_stats_count, src_ip, dst_ip, header->len);
}

int main() {
    char errbuf[PCAP_ERRBUF_SIZE];
    pcap_if_t* all_devs;
    pcap_t* handle;
    const char* dev; // 选择捕获的网络接口

    // 获取所有设备
    if (pcap_findalldevs(&all_devs, errbuf) == -1) {
        fprintf(stderr, "Error finding devices: %s\n", errbuf);
        return 1;
    }

    // 使用第一个设备
    dev = all_devs->name;
    printf("Capturing on device: %s\n", dev);

    // 打开设备进行捕获
    handle = pcap_open_live(dev, BUFSIZ, 1, 1000, errbuf);
    if (handle == NULL) {
        fprintf(stderr, "Couldn't open device %s: %s\n", dev, errbuf);
        return 2;
    }

    // 创建 CSV 文件并写入标题
    FILE* log_file = fopen(CSV_LOG_FILE, "w");
    if (log_file != NULL) {
        fprintf(log_file, "时间,源MAC,源IP,目标MAC,目标IP,帧长度\n");
        fclose(log_file);
    }

    time_t start_time = time(NULL);

    // 捕获数据包
    while (1) {
        pcap_loop(handle, 1, packet_handler, NULL);

        // 每隔 INTERVAL 秒记录一次统计信息
        if (time(NULL) - start_time >= INTERVAL) {
            log_stats_to_csv();
            start_time = time(NULL);
        }
    }

    pcap_close(handle);
    return 0;
}
